// Copyright 2024 TiKV Project Authors. Licensed under Apache-2.0.
use std::{
    path::{Path, PathBuf},
    sync::Arc,
};

use engine_rocks::RocksEngine;
use engine_traits::{
    ExternalSstFileInfo, SstCompressionType, SstExt, SstWriter, SstWriterBuilder,
    DATA_KEY_PREFIX_LEN,
};
use external_storage::{ExternalStorage, UnpinReader};
use file_system::Sha256Reader;
use futures::{future::TryFutureExt, io::AllowStdIo};
use kvproto::brpb::{self, LogFileSubcompaction};
use tikv_util::{
    codec::bytes::decode_bytes_in_place, retry_expr, stream::JustRetry, time::Instant,
};

use super::{EpochHint, Subcompaction, SubcompactionResult};
use crate::{
    compaction::SST_OUT_REL,
    errors::{OtherErrExt, Result, TraceResultExt},
    source::{Record, Source},
    statistic::{prom::*, LoadStatistic, SubcompactStatistic},
    storage::DEFAULT_COMPACTION_OUT_PREFIX,
    util::{self, Cooperate, ExecuteAllExt},
};

/// The state of executing a subcompaction.
pub struct SubcompactionExec<DB> {
    source: Source,
    output: Arc<dyn ExternalStorage>,
    co: Cooperate,
    out_prefix: PathBuf,

    load_stat: LoadStatistic,
    compact_stat: SubcompactStatistic,

    db: Option<DB>,
}

/// The extra config of executing a subcompaction.
pub struct SubcompactExt {
    /// The number of max concurrency of loading input files.
    pub max_load_concurrency: usize,
    /// The compression type that will be used for the output file.
    pub compression: SstCompressionType,
    /// The compression level of the output file.
    ///
    /// When `None`, will use the default level of the compression algorithm.
    pub compression_level: Option<i32>,
}

impl Default for SubcompactExt {
    fn default() -> Self {
        Self {
            max_load_concurrency: Default::default(),
            compression: SstCompressionType::Lz4,
            compression_level: None,
        }
    }
}

// NOTE: maybe we can merge this and `SubcompactionExt`?
/// The information used for constructing a [`SubcompactionExec`].
pub struct SubcompactionExecArg<DB> {
    /// The prefix of the output SST.
    pub out_prefix: Option<PathBuf>,
    /// The RocksDB instance used for creating the SST writer.
    pub db: Option<DB>,
    /// The output storage.
    pub storage: Arc<dyn ExternalStorage>,
}

impl<DB> From<SubcompactionExecArg<DB>> for SubcompactionExec<DB> {
    fn from(value: SubcompactionExecArg<DB>) -> Self {
        Self {
            source: Source::new(Arc::clone(&value.storage)),
            output: value.storage,
            out_prefix: value
                .out_prefix
                .unwrap_or_else(|| Path::new(DEFAULT_COMPACTION_OUT_PREFIX).to_owned()),
            db: value.db,

            co: Default::default(),
            load_stat: Default::default(),
            compact_stat: Default::default(),
        }
    }
}

impl SubcompactionExec<RocksEngine> {
    #[cfg(test)]
    pub fn default_config(storage: Arc<dyn ExternalStorage>) -> Self {
        Self::from(SubcompactionExecArg {
            storage,
            out_prefix: None,
            db: None,
        })
    }
}

/// An SST generated by a subcompaction.
struct WrittenSst<S> {
    content: S,
    meta: kvproto::brpb::File,
    epoch_hints: Vec<EpochHint>,
    physical_size: u64,
}

/// Log backup may generate duplicated key-value entries.
/// When compacting them in the same subcompaction, we will dedup them. Hence
/// the checksum of output varies.
///
/// During compacting, those difference will be recorded here.
#[derive(Default)]
struct ChecksumDiff {
    removed_key: u64,
    decreaed_size: u64,
    crc64xor_diff: u64,
}

impl<DB: SstExt> SubcompactionExec<DB>
where
    <<DB as SstExt>::SstWriter as SstWriter>::ExternalSstFileReader: 'static,
{
    fn update_checksum_diff(a: &Record, b: &Record, diff: &mut ChecksumDiff) {
        assert_eq!(
            a, b,
            "The record with same key contains different value: the backup might be corrupted."
        );

        diff.removed_key += 1;
        diff.decreaed_size += (a.key.len() + a.value.len()) as u64;
        let mut d = crc64fast::Digest::new();
        d.write(&a.key);
        d.write(&a.value);
        // When we remove even number of the same key-value pair, we don't need to
        // update the crc64xor checksum. Trick here: If a key-value pair appears
        // even times, that key will be eliminated in this checksum, so we just
        // need to `XOR` all duplicated keys up, we will get the right crc64xor diff.
        diff.crc64xor_diff ^= d.sum64();
    }

    /// Sort all inputs and dedup them, generating the checksum diff.
    #[tracing::instrument(skip_all)]
    async fn process_input(
        &mut self,
        items: impl Iterator<Item = Vec<Record>>,
    ) -> (Vec<Record>, ChecksumDiff) {
        let mut flatten_items = items
            .into_iter()
            .flat_map(|v| v.into_iter())
            .collect::<Vec<_>>();
        tokio::task::yield_now().await;
        flatten_items.sort_unstable_by(Record::cmp_key);
        tokio::task::yield_now().await;
        let mut diff = ChecksumDiff::default();
        flatten_items.dedup_by(|k1, k2| {
            if k1.key == k2.key {
                Self::update_checksum_diff(k1, k2, &mut diff);
                true
            } else {
                false
            }
        });
        (flatten_items, diff)
    }

    #[tracing::instrument(skip_all)]
    async fn load(
        &mut self,
        c: &Subcompaction,
        ext: &mut SubcompactExt,
    ) -> Result<impl Iterator<Item = Vec<Record>>> {
        let mut eext = ExecuteAllExt::default();
        eext.max_concurrency = ext.max_load_concurrency;

        let items = super::util::execute_all_ext(
            c.inputs
                .iter()
                .cloned()
                .map(|f| {
                    let source = &self.source;
                    Box::pin(async move {
                        let _lt = COMPACT_LOG_BACKUP_LOAD_A_FILE_DURATION.start_coarse_timer();
                        let mut out = vec![];
                        let mut stat = LoadStatistic::default();
                        source
                            .load(f, Some(&mut stat), |k, v| {
                                fail::fail_point!("compact_log_backup_omit_key", |_| {});
                                out.push(Record {
                                    key: k.to_owned(),
                                    value: v.to_owned(),
                                })
                            })
                            .await?;
                        Result::Ok((out, stat))
                    })
                })
                .collect(),
            eext,
        )
        .await?;

        let mut result = Vec::with_capacity(items.len());
        for (item, stat) in items {
            self.load_stat += stat;
            result.push(item);
        }

        Ok(result.into_iter())
    }

    /// write the `sorted_items` to a in-mem SST.
    ///
    /// # Panics
    ///
    /// For now, if the `sorted_items` is empty, it will panic.
    /// But it is reasonable to return an error in this scenario if needed.
    #[tracing::instrument(skip_all, fields(name=%name))]
    async fn write_sst(
        &mut self,
        name: &str,
        c: &Subcompaction,
        sorted_items: &[Record],
        ext: &mut SubcompactExt,
    ) -> Result<WrittenSst<<DB::SstWriter as SstWriter>::ExternalSstFileReader>> {
        let cf = c.cf;
        let mut wb = <DB as SstExt>::SstWriterBuilder::new()
            .set_cf(cf)
            .set_compression_type(Some(ext.compression))
            .set_in_memory(true);
        if let Some(db) = self.db.as_ref() {
            wb = wb.set_db(db);
        }
        if let Some(level) = ext.compression_level {
            wb = wb.set_compression_level(level);
        }
        let mut w = wb.build(name)?;
        let mut meta = kvproto::brpb::File::default();

        let mut start_key = sorted_items[0].key.clone();
        // `File::{start,end}_key` should be raw key.
        decode_bytes_in_place(&mut start_key, false).adapt_err()?;
        let mut end_key = sorted_items.last().unwrap().key.clone();
        decode_bytes_in_place(&mut end_key, false).adapt_err()?;
        // `File::end_key` should be exclusive. (!)
        // Also we cannot just call next_key, or the table ID of the end key may be
        // different, some versions of BR panics in that scenario.
        end_key.push(0u8);

        meta.set_start_key(start_key);
        meta.set_end_key(end_key);
        meta.set_cf(cf.to_owned());
        meta.name = name.to_owned();
        meta.end_version = u64::MAX;

        let mut data_key = keys::DATA_PREFIX_KEY.to_vec();
        for item in sorted_items {
            self.co.step().await;

            let mut d = crc64fast::Digest::new();
            d.write(&item.key);
            d.write(&item.value);
            let ts = item.ts().trace_err()?;
            meta.crc64xor ^= d.sum64();
            meta.start_version = meta.start_version.min(ts);
            meta.end_version = meta.end_version.max(ts);

            // NOTE: We may need to check whether the key is already a data key here once we
            // are going to support compact SSTs.
            data_key.truncate(DATA_KEY_PREFIX_LEN);
            data_key.extend(&item.key);
            w.put(&data_key, &item.value)?;

            self.compact_stat.logical_key_bytes_out += item.key.len() as u64;
            self.compact_stat.logical_value_bytes_out += item.value.len() as u64;
            meta.total_kvs += 1;
            meta.total_bytes += item.key.len() as u64 + item.value.len() as u64;
        }
        let (info, out) = w.finish_read()?;
        meta.set_size(info.file_size());
        self.compact_stat.keys_out += info.num_entries();
        self.compact_stat.physical_bytes_out += info.file_size();

        let result = WrittenSst {
            content: out,
            meta,
            epoch_hints: c.epoch_hints.clone(),
            physical_size: info.file_size(),
        };

        Ok(result)
    }

    #[tracing::instrument(skip_all, fields(name=%sst.meta.name))]
    async fn upload_compaction_artifact(
        &mut self,
        c: &Subcompaction,
        sst: &mut WrittenSst<<DB::SstWriter as SstWriter>::ExternalSstFileReader>,
    ) -> Result<LogFileSubcompaction> {
        use engine_traits::ExternalSstFileReader;
        sst.content.reset()?;
        let (rd, hasher) = Sha256Reader::new(&mut sst.content).adapt_err()?;
        self.output
            .write(
                &sst.meta.name,
                UnpinReader(Box::new(AllowStdIo::new(rd))),
                sst.physical_size,
            )
            .await?;
        sst.meta.sha256 = hasher.lock().unwrap().finish().adapt_err()?.to_vec();
        let mut meta = brpb::LogFileSubcompaction::new();
        meta.set_meta(c.to_pb_meta());
        meta.set_region_meta_hints(sst.epoch_hints.iter().map(|v| v.to_pb()).collect());
        meta.set_sst_outputs(vec![sst.meta.clone()].into());
        Ok(meta)
    }

    #[tracing::instrument(skip_all, fields(c=%c))]
    pub async fn run(
        mut self,
        c: Subcompaction,
        mut ext: SubcompactExt,
    ) -> Result<SubcompactionResult> {
        let mut result = SubcompactionResult::of(c);
        let c = &result.origin;
        for input in &c.inputs {
            if input.crc64xor == 0 {
                result.expected_crc64 = None;
            }
            if let Some(v) = result.expected_crc64.as_mut() {
                *v ^= input.crc64xor;
            }
            result.expected_keys += input.num_of_entries;
            result.expected_size += input.key_value_size;
        }

        let begin = Instant::now();
        let items = self.load(c, &mut ext).await?;
        self.compact_stat.load_duration += begin.saturating_elapsed();

        let begin = Instant::now();
        let (sorted_items, cdiff) = self.process_input(items).await;
        self.compact_stat.sort_duration += begin.saturating_elapsed();
        if sorted_items.is_empty() {
            self.compact_stat.empty_generation += 1;
            return Ok(result);
        }

        if let Some(v) = result.expected_crc64.as_mut() {
            *v ^= cdiff.crc64xor_diff;
        }
        result.expected_keys -= cdiff.removed_key;
        result.expected_size -= cdiff.decreaed_size;

        let out_name = self
            .out_prefix
            .join(SST_OUT_REL)
            .join(format!(
                "{}_{}_{}_{}.sst",
                util::aligned_u64(c.input_min_ts),
                util::aligned_u64(c.input_max_ts),
                c.cf,
                c.region_id
            ))
            .display()
            .to_string();
        let begin = Instant::now();
        assert!(!sorted_items.is_empty());
        let mut sst = self
            .write_sst(&out_name, c, sorted_items.as_slice(), &mut ext)
            .await?;

        self.compact_stat.write_sst_duration += begin.saturating_elapsed();

        let begin = Instant::now();
        result.meta =
            retry_expr! { self.upload_compaction_artifact(c, &mut sst).map_err(JustRetry) }
                .await
                .map_err(|err| err.0)?;
        self.compact_stat.save_duration += begin.saturating_elapsed();

        result.compact_stat = self.compact_stat;
        result.load_stat = self.load_stat;

        return Ok(result);
    }
}

#[cfg(test)]
mod test {
    use tidb_query_datatype::codec::table::encode_row_key;
    use txn_types::Key;

    use crate::{
        compaction::Subcompaction,
        storage::{Epoch, MetaFile},
        test_util::{
            gen_step, save_many_log_files, CompactInMem, KvGen, LogFileBuilder, TmpStorage,
        },
    };

    #[tokio::test]
    async fn test_compact_one() {
        let st = TmpStorage::create();

        let const_val = |_| vec![42u8];
        let cm = CompactInMem::default();

        let s1 = KvGen::new(gen_step(1, 0, 2).take(100), const_val);
        let i1 = st.build_log_file("a.log", cm.tap_on(s1)).await;

        let s2 = KvGen::new(gen_step(1, 1, 2).take(100), const_val);
        let i2 = st.build_log_file("b.log", cm.tap_on(s2)).await;

        let c = Subcompaction::of_many([i1, i2]);

        st.verify_result(st.run_subcompaction(c).await, cm);
    }

    #[tokio::test]
    async fn test_compact_dup() {
        let st = TmpStorage::create();
        let cm = CompactInMem::default();

        let s1 = KvGen::new(gen_step(1, 0, 3).take(100), |_| b"value".to_vec());
        let i1 = st.build_log_file("three.log", cm.tap_on(s1)).await;

        let s2 = KvGen::new(gen_step(1, 0, 2).take(100), |_| b"value".to_vec());
        let i2 = st.build_log_file("two.log", cm.tap_on(s2)).await;

        let c = Subcompaction::of_many([i1, i2]);
        let res = st.run_subcompaction(c).await;
        assert_eq!(res.load_stat.keys_in, 200);
        assert_eq!(res.compact_stat.keys_out, 166);
        st.verify_result(res, cm);
    }

    #[tokio::test]
    async fn test_compact_from_one_file() {
        let st = TmpStorage::create();
        let cm = CompactInMem::default();

        let s1 = KvGen::new(gen_step(1, 0, 2).take(100), |_| b"value".to_vec());
        let mut i1 = LogFileBuilder::new(|v| v.name = "a.log".to_owned());
        cm.tap_on(s1)
            .for_each(|kv| i1.add_encoded(&kv.key, &kv.value));

        let s2 = KvGen::new(gen_step(1, 1, 2).take(128), |_| b"value".to_vec());
        let mut i2 = LogFileBuilder::new(|v| v.name = "b.log".to_owned());
        cm.tap_on(s2)
            .for_each(|kv| i2.add_encoded(&kv.key, &kv.value));

        let meta = save_many_log_files("data.log", [i1, i2], st.storage().as_ref())
            .await
            .unwrap();
        let ml = MetaFile::from(meta);
        let c = Subcompaction::of_many(ml.into_logs());
        let res = st.run_subcompaction(c).await;
        st.verify_result(res, cm);
    }

    #[tokio::test]
    async fn test_region_boundaries() {
        let st = TmpStorage::create();
        let cm = CompactInMem::default();

        let s1 = KvGen::new(gen_step(1, 0, 2).take(100), |_| b"value".to_vec());
        let enc = |v| Key::from_raw(v).into_encoded();
        let mut i1 = LogFileBuilder::new(|v| {
            v.name = "a.log".to_owned();
            v.region_start_key = Some(enc(b"t"));
            v.region_end_key = Some(enc(b"t\xff"));
            v.region_epoches.push(Epoch {
                version: 42,
                conf_ver: 42,
            });
        });
        cm.tap_on(s1)
            .for_each(|kv| i1.add_encoded(&kv.key, &kv.value));
        let meta = save_many_log_files("data.log", [i1], st.storage().as_ref())
            .await
            .unwrap();
        let ml = MetaFile::from(meta);
        let c = Subcompaction::of_many(ml.into_logs());
        let res = st.run_subcompaction(c).await;
        let sst_out = &res.meta.get_region_meta_hints()[0];
        assert_eq!(sst_out.get_start_key(), enc(b"t").as_slice());
        assert_eq!(sst_out.get_end_key(), enc(b"t\xff").as_slice());
        st.verify_result(res, cm);
    }

    #[tokio::test]
    async fn test_elide_region_boundaries() {
        let st = TmpStorage::create();
        let cm = CompactInMem::default();

        let s1 = KvGen::new(gen_step(1, 0, 2).take(100), |_| b"value".to_vec());
        let enc = |v| Key::from_raw(v).into_encoded();
        let mut i1 = LogFileBuilder::new(|v| {
            v.name = "a.log".to_owned();
            v.region_start_key = Some(enc(b"t"));
            v.region_end_key = Some(enc(b"t\xff"));
            v.region_epoches.push(Epoch {
                version: 42,
                conf_ver: 42,
            });
        });
        cm.tap_on(s1)
            .for_each(|kv| i1.add_encoded(&kv.key, &kv.value));

        let s2 = KvGen::new(gen_step(1, 200, 2).take(100), |_| b"value".to_vec());
        let mut i2 = LogFileBuilder::new(|v| {
            v.name = "b.log".to_owned();
        });
        cm.tap_on(s2)
            .for_each(|kv| i2.add_encoded(&kv.key, &kv.value));

        let meta = save_many_log_files("data.log", [i1, i2], st.storage().as_ref())
            .await
            .unwrap();
        let ml = MetaFile::from(meta);
        let logs = ml.into_logs().collect::<Vec<_>>();
        assert_eq!(logs.len(), 2, "{:?}", logs);

        // Case 1: with ranges only.
        let c = Subcompaction::of_many([logs[0].clone()]);
        let res = st.run_subcompaction(c).await;
        let sst_out = &res.meta.get_region_meta_hints()[0];
        let bgn = sst_out.get_start_key();
        let end = sst_out.get_end_key();
        let cbgn = res.meta.get_meta().get_min_key();
        let cend = res.meta.get_meta().get_max_key();
        assert_eq!(bgn, enc(b"t").as_slice());
        assert_eq!(end, enc(b"t\xff").as_slice());
        assert!(cbgn > enc(bgn).as_slice(), "{:?}", cbgn);
        assert!(cend < enc(end).as_slice(), "{:?}", cend);

        // Case 2: should elide range.
        let c = Subcompaction::of_many(logs);
        let res = st.run_subcompaction(c).await;
        let sst_out = &res.meta.get_sst_outputs()[0];
        let bgn = sst_out.get_start_key();
        let end = sst_out.get_end_key();
        let cbgn = res.meta.get_meta().get_min_key();
        let cend = res.meta.get_meta().get_max_key();
        let row = |r| encode_row_key(1, r);
        assert_eq!(bgn, row(0));
        let mut ek = row(398);
        ek.push(0);
        assert_eq!(end, ek);
        assert!(cbgn > enc(bgn).as_slice(), "{:?}", cbgn);
        assert!(cend < enc(end).as_slice(), "{:?}", cend);
        st.verify_result(res.clone(), cm);
    }

    #[tokio::test]
    #[cfg(feature = "failpoints")]
    // Note: maybe just modify the log files?
    async fn test_failed_checksumming() {
        let _fg = fail::FailGuard::new("compact_log_backup_omit_key", "1*return");
        let st = TmpStorage::create();

        let cm = CompactInMem::default();

        let s1 = KvGen::new(gen_step(1, 0, 3).take(100), |_| b"value".to_vec());
        let i1 = st.build_log_file("three.log", cm.tap_on(s1)).await;

        let c = Subcompaction::singleton(i1);
        let res = st.run_subcompaction(c).await;
        res.verify_checksum()
            .expect_err("should failed to verify checksum");
    }
}
